/*
 * kmean.h
 *
 *  Created on: Mar 6, 2011
 *      Author: tqlong
 */

#ifndef KMEAN_H_
#define KMEAN_H_

#include <cstdlib>
#include <limits>

/** Classical KMean clustering algorithm */
template<typename VEC, typename DATASET>
class KMean
{
public:
  typedef VEC                      vec_type;
  typedef DATASET                  dataset_type;
  typedef std::vector<double>      dvec_type;
  typedef std::vector<dvec_type>   vec_dvec_type;
protected:
  const dataset_type& data_;                       // the data
  int n_class_;                                    // number of classes
  int n_sample_;                                   // number of samples
  int n_dim_;                                      // dimensionality
  std::vector<vec_type> center_;                   // centers of the clusters
  dvec_type nc_;                                   // number of samples in each class
  dvec_type y_;                                    // class labels
  int maxiter_;                                    // maximum number of iterations
public:
  KMean(const dataset_type& data, int K)
    : data_(data), n_class_(K), n_sample_(data.n()), n_dim_(data.dim()),
      center_(K, vec_type(n_dim_)), nc_(K), y_(n_sample_), maxiter_(100)
  {
  }

  void setParams(const std::vector<double>& p)
  {
    if (p.size()>0) maxiter_ = (int)p[0];
  }

  void initCenters()
  {
    center_[0] = data_[::rand() % n_sample_];              // init the centers
    for (int k = 1; k < n_class_; k++) {                   // by choosing n_class centers
      double maxd = -1.0;                                  // that are very far from each other
      int maxi = -1;
      for (int i = 0; i < n_sample_; i++) {
        double mind = arma::norm(data_[i]-center_[0], 2);
        for (int k1 = 0; k1 < k; k1++) {
          double d = arma::norm(data_[i]-center_[k1], 2);
          if (d < mind) {
            mind = d;
          }
        }
        if (mind > maxd) {
          maxd = mind;
          maxi = i;
        }
      }
      assert(maxi >= 0);
      center_[k] = data_[maxi];
    }

    std::cerr << "kmean -- iter = " << 0 << " --\n";
//    for (int i = 0; i < n_class_; i++) {
//      std::cerr << "center[" << i << "]=\n" << center_[i] << "\n";
//    }
  }

  void optimize()
  {
    bool no_change;
    initCenters();
    assignClusters(no_change);
    no_change = false;
    int iter = 1;
    while (!no_change && iter < maxiter_) {                          // repeat computing centers
      computeCenters();                                              // and reassign points to clusters
      std::cerr << "-- kmean iter = " << iter << " --\n";                  // until there are no change
      for (int i = 0; i < n_class_; i++)
        std::cerr << "center[" << i << "]=\n" << center_[i] << "\n";
      assignClusters(no_change);
      iter++;
    }
  }

  void assignClusters(bool &no_change)
  {
    no_change = true;                                     // assign each point to its
    for (int i = 0; i < n_sample_; i++) {                 // closest center
      int y_old = (int)y_[i];
      int mink = 0;
      double mind = arma::norm(data_[i]-center_[0], 2);
      for (int k = 1; k < n_class_; k++) {
        double d = arma::norm(data_[i]-center_[k], 2);
        if (d < mind) {
          mind = d;
          mink = k;
        }
      }
      if (mink != y_old) no_change = false;
      y_[i] = mink;
    }
  }

  void computeCenters()
  {
    for (int k = 0; k < n_class_; k++) {              // computer new center by averaging
      center_[k] = vec_type(n_dim_);                  // points assigned to the same cluster
      center_[k].zeros();
      nc_[k] = 0.0;
      for (int i = 0; i < n_sample_; i++)
        if ((int)y_[i] == k) {
          nc_[k] += 1.0;
          center_[k] += data_[i];
        }
      if (nc_[k] > 0)
        center_[k] /= nc_[k];
    }
    for (int k = 0; k < n_class_; k++)                // exception are "bad" clusters with no point
      if (nc_[k] < 1.0) {                             // randomly re-initialize theirs centers
        double maxd = -1.0;                           // to points that are far from "good" centers
        int maxi = -1;
        for (int i = 0; i < n_sample_; i++) {
          double mind = std::numeric_limits<double>::infinity();
          for (int k1 = 0; k1 < n_class_; k1++) if (nc_[k1]>0) {
            double d = arma::norm(data_[i]-center_[k1], 2);
            if (d < mind) {
              mind = d;
            }
          }
          if (mind > maxd) {
            maxd = mind;
            maxi = i;
          }
        }
        assert(maxi >= 0);
        center_[k] = data_[maxi];
        nc_[k] = 1;
      }
  }

  double y(int i) const { return y_[i]; }
  const vec_type& center(int k) const { return center_[k]; }
};

#endif /* KMEAN_H_ */
